#!/usr/bin/env python
#-*- coding: UTF-8 -*-
# nodif.py: node manager interface functions for node client
#
# Copyright (C) 2008-2010 Hily Jiang <hilyjiang@gmail.com>
#
# This file is part of dtask.
#
# This software is provided 'as-is', without any express or implied
# warranty.  In no event will the authors be held liable for any damages
# arising from the use of this software.
#
# Permission is granted to anyone to use this software for any purpose,
# including commercial applications, and to alter it and redistribute it
# freely, subject to the following restrictions:
#
# 1. The origin of this software must not be misrepresented; you must not
#    claim that you wrote the original software. If you use this software
#    in a product, an acknowledgment in the product documentation would be
#    appreciated but is not required.
# 2. Altered source versions must be plainly marked as such, and must not be
#    misrepresented as being the original software.
# 3. This notice may not be removed or altered from any source distribution.

"""
Node Manager interface functions for Node Client.

Node <-> Node Manager Protocol:

* New Connection
  (N -> NM) Send data format:    PASSPORT nodType nodName nodKey\n

* Close Connection
  (N -> NM) Send data format:    DISCONNECT\n

* Dispatch Task
  (NM -> N) Send data format:    TASK taskId taskType taskParams\n

* Send Result
  (N -> NM) Send data format:    RESULT taskId result\n
"""

import sys
import threading

from utils import CommandBuffer


CFG_NOD_BUFSIZ = 1024


class NodIfBuffer(CommandBuffer):
    """Node buffer to receive command string and node info."""

    def __init__(self):
        CommandBuffer.__init__(self)
        self.nodName = ''
        self.nodKey = ''
        self.hasInfo = False
        self.nodeId = None

    def setInfo(self, nodName='', nodKey=''):
        """Set node info."""
        if self.hasInfo:
            return False
        self.nodName = nodName
        self.nodKey = nodKey
        self.hasInfo = True

    def getInfo(self):
        """Get node info."""
        return (self.nodName, self.nodKey)

    def setNodeId(self, nodeId):
        self.nodeId = nodeId

    def getNodeId(self):
        return self.nodeId


def nodCommandHandler(cmdInfo, nodeManager, dispatcher, niBuf,
                      nodClientSock, nodClientName,
                      nodSockDict, schSockDict):
    """Handle the node's commands."""
    cmd, params = cmdInfo

    print nodClientName, '[NOD]', sys._getframe().f_lineno, cmd, 
    if params:
        print params[:40]
    else:
        print
    if cmd == 'RESULT':
        nodeId = niBuf.getNodeId()
        nodeManager.releaseNode(nodeId)
        try:
            ps = params.split(None, 1)
            if len(ps) == 2:
                taskId, result = ps
            else:
                taskId = ps[0]
                result = ''
            schId, taskId = taskId.split('|', 1) # 'schid|orginalTaskId'
        except Exception, e:
            print nodClientName, '[NOD]', sys._getframe().f_lineno, e
        else:
            nodIp, nodPort = nodClientName
            schId = dispatcher.getSchedulerId(schId, taskId, nodeId)
            if schId and schSockDict.has_key(schId):
                schSockDict[schId].sendall('RESULT %s %s\n' % (taskId, result))
    elif cmd == 'PASSPORT':
        params = params.split(None, 2)
        if len(params) == 1:
            nodType = params[0]
            nodName = nodKey = ''
        elif len(params) == 2:
            nodType, nodName = params
            nodKey = ''
        else:
            nodType, nodName, nodKey = params
        niBuf.setInfo(nodName, nodKey)
        nodIp, nodPort = nodClientName
        nodeId = nodeManager.newNode(nodIp, nodPort, nodType, nodName)
        nodSockDict[nodeId] = nodClientSock
        niBuf.setNodeId(nodeId)
        # todo: authorization
    elif cmd == 'DISCONNECT':
        nodeId = niBuf.getNodeId()
        nodeManager.removeNode(nodeId)
        if nodSockDict.has_key(nodeId):
            del nodSockDict[nodeId]


def nodRecvData(nodeManager, dispatcher, niBuf,
                nodClientSock, nodClientName, data,
                nodSockDict, schSockDict):
    """Process data send by node."""
    niBuf.addData(data)

    while True:
        cmdInfo = niBuf.getCommand()
        if cmdInfo:
            nodCommandHandler(cmdInfo, nodeManager, dispatcher, niBuf,
                              nodClientSock, nodClientName,
                              nodSockDict, schSockDict)
        else:
            break


def nodRecvThreadFunc(nodeManager, dispatcher, nodClientSock, nodClientSocks,
                      nodSockDict, schSockDict):
    """Thread function for each node connection to process data."""
    nodClientName = nodClientSock.getpeername()
    niBuf = NodIfBuffer()

    while True:
        try:
            data = nodClientSock.recv(CFG_NOD_BUFSIZ)
            if data == '': break
            nodRecvData(nodeManager, dispatcher, niBuf,
                        nodClientSock, nodClientName, data,
                        nodSockDict, schSockDict)
        except Exception, e:
            print nodClientName, '[NOD]', sys._getframe().f_lineno, e
    
    print nodClientName, '[NOD] quit'
    nodeId = niBuf.getNodeId()
    nodeManager.removeNode(nodeId)
    if nodSockDict.has_key(nodeId):
        del nodSockDict[nodeId]
    nodClientSock.close()
    nodClientSocks.remove(nodClientSock)


def startNodeIf(nodeManager, dispatcher, nodServerSock, nodClientSocks,
                nodSockDict, schSockDict):
    """Start interface for nodes."""
    nodRecvThreads = []

    while True:
        nodClientSock, nodClientAddr = nodServerSock.accept()
        nodClientSocks.append(nodClientSock)
        print nodClientAddr, '[NOD] connected'
        t = threading.Thread(target=nodRecvThreadFunc,
                   args=(nodeManager, dispatcher, nodClientSock,
                         nodClientSocks, nodSockDict, schSockDict))
        t.setDaemon(True)
        nodRecvThreads.append(t)
        t.start()

    print 'Waiting for node clients to release...',
    for t in nodRecvThreads:
        t.join()
    print 'ok'
